import cv2
import math
import numpy as np
from PyQt5.QtWidgets import QWidget, QLabel, QVBoxLayout
from PyQt5.QtCore import QTimer
from PyQt5.QtGui import QImage, QPixmap
import mediapipe as mp
import time

class FaceWidget(QWidget):
    def __init__(self):
        super().__init__()
        print("Initializing FaceWidget...")
        self.label = QLabel("Camera Loading...")
        layout = QVBoxLayout()
        layout.addWidget(self.label)
        self.setLayout(layout)

        # 初始化摄像头
        self.cap = None
        self.init_camera()
        
        self.mp_face_detection = mp.solutions.face_detection
        self.mp_drawing = mp.solutions.drawing_utils
        self.face_detection = self.mp_face_detection.FaceDetection(min_detection_confidence=0.8)

        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frame)
        self.timer.start(100)  # 降低帧率到10fps
        print("FaceWidget initialization complete")

    def init_camera(self):
        """初始化摄像头，尝试不同的摄像头索引"""
        for i in range(2):  # 尝试摄像头0和1
            print(f"Trying camera {i}...")
            self.cap = cv2.VideoCapture(i)
            if self.cap.isOpened():
                # 设置摄像头分辨率
                self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
                self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
                # 等待摄像头初始化
                time.sleep(1)
                ret, _ = self.cap.read()
                if ret:
                    print(f"Camera {i} opened successfully")
                    return
                else:
                    self.cap.release()
        
        print("Error: Could not open any camera")
        self.label.setText("Error: Camera not available")

    def closeEvent(self, event):
        print("Closing FaceWidget...")
        if self.cap is not None:
            self.cap.release()
        event.accept()

    def update_frame(self):
        if self.cap is None or not self.cap.isOpened():
            print("Camera not available, trying to reinitialize...")
            self.init_camera()
            return

        try:
            ret, img = self.cap.read()
            if not ret:
                print("Error: Could not read frame, trying to reinitialize camera...")
                self.cap.release()
                self.init_camera()
                return

            frame = self.process_frame(img)
            self.display_image(frame)
        except Exception as e:
            print(f"Error in update_frame: {str(e)}")
            # 发生错误时尝试重新初始化摄像头
            if self.cap is not None:
                self.cap.release()
            self.init_camera()

    def process_frame(self, img):
        try:
            img.flags.writeable = False
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            results = self.face_detection.process(img_rgb)
            img.flags.writeable = True
            h, w = img.shape[:2]

            direction_text = "No Face Detected"

            if results.detections:
                for detection in results.detections:
                    self.mp_drawing.draw_detection(img, detection)
                    data = detection.location_data.relative_bounding_box
                    point = self._normalized_to_pixel_coordinates(data.xmin, data.ymin, w, h)
                    if point:
                        x1, y1 = point
                        width, height = self._normalized_to_pixel_coordinates(data.width, data.height, w, h)
                        x, y = (x1 + width / 2, y1 + height / 2)
                        cv2.circle(img, (int(x), int(y)), 5, (255, 255, 0), -1)
                        cv2.putText(img, f"({int(x)}, {int(y)})", (int(x)-50, int(y)-15),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 0), 2)

                        center_x = w / 2
                        tolerance = w / 10
                        if x < center_x - tolerance:
                            direction_text = "Turn Left"
                        elif x > center_x + tolerance:
                            direction_text = "Turn Right"
                        else:
                            direction_text = "Face Centered"

            cv2.putText(img, direction_text, (30, 50), cv2.FONT_HERSHEY_SIMPLEX,
                        1, (0, 255, 255), 2)

            return img
        except Exception as e:
            print(f"Error in process_frame: {str(e)}")
            return img

    def display_image(self, img):
        try:
            rgb_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            bytes_per_line = ch * w
            qt_img = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
            pix = QPixmap.fromImage(qt_img)
            self.label.setPixmap(pix)
        except Exception as e:
            print(f"Error in display_image: {str(e)}")

    def _normalized_to_pixel_coordinates(self, normalized_x, normalized_y, image_width, image_height):
        def is_valid(value): return 0 <= value <= 1
        if not (is_valid(normalized_x) and is_valid(normalized_y)):
            return None
        x_px = min(math.floor(normalized_x * image_width), image_width - 1)
        y_px = min(math.floor(normalized_y * image_height), image_height - 1)
        return x_px, y_px

    def start_camera(self):
        """启动摄像头"""
        if self.cap is None or not self.cap.isOpened():
            self.init_camera()
        self.timer.start(100)

    def stop_camera(self):
        """停止摄像头"""
        self.timer.stop()
        if self.cap is not None:
            self.cap.release()
            self.cap = None

    def reinit_camera(self):
        """重新初始化摄像头"""
        self.stop_camera()
        self.start_camera()
